#! /usr/bin/env python
# -*- coding: utf-8 -*-

# Copyright (c) Huoty, All rights reserved
# Author: Huoty <sudohuoty@163.com>
# CreateTime: 2022-05-18 19:08:52

import base64
import hashlib


def rc4(text, key='default', mode="encode"):
    # 不直接使用原始秘钥，而是取长度固定为 32 的 MD5 摘要值作为密钥，从而方便处理
    key = hashlib.md5(key.encode("utf8")).hexdigest()

    # 加密得到的消息使用 base64 编码，则解密的时候需要先将 base64 编码去除
    if mode == "decode":
        text = base64.b64decode(text).decode("utf8")

    result = ''
    key_len = len(key)

    # init S-box
    box = list(range(256))
    j = 0
    for i in range(256):
        j = (j + box[i] + ord(key[i % key_len])) % 256
        box[i], box[j] = box[j], box[i]

    # make sure all elements in S-box swapped at least once
    i = j = 0
    for element in text:
        i = (i + 1) % 256
        j = (j + box[i]) % 256
        box[i], box[j] = box[j], box[i]
        k = chr(ord(element) ^ box[(box[i] + box[j]) % 256])
        result += k

    if mode == "encode":
        result = base64.b64encode(result.encode("utf8")).decode("utf8")

    return result


if __name__ == "__main__":
    key = "ABCL60JsdRQc7lB4ayK9DyM1n1X8YUjV"
    crypt_text = rc4("hello world", key=key)
    print(rc4(crypt_text, key=key, mode="decode"), "\n", crypt_text)
